package org.zcl.elastic.dao;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Sort;
import org.zcl.elastic.entity.Scroll;

import java.io.IOException;
import java.util.Date;
import java.util.Iterator;
import java.util.List;

/**
 * .
 *
 * @author zhuchenglin
 * @since 2024/9/8 21:19
 */
@Slf4j
public class EsRepository {

    @Autowired
    private RestHighLevelClient client;

    @Autowired
    private ObjectMapper objectMapper;

    public <T> Scroll<T> findScroll(String scrollId, Class<T> objectClass) {
        SearchScrollRequest request = new SearchScrollRequest(scrollId);
        request.scroll(TimeValue.timeValueSeconds(60L));

        SearchResponse searchResponse;
        try {
            searchResponse = this.client.scroll(request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            throw new RuntimeException("Find scroll document from elasticSearch by scrollId error.", e);
        }

        SearchHits hits = searchResponse.getHits();
        List<T> content = this.parseHits(hits, objectClass);
        return new Scroll(scrollId, new Date((new Date()).getTime() + 60000L), hits.getTotalHits().value, content);
    }


    public <T> Scroll<T> findScroll(String index, String[] fields, int scrollSize, Sort sort, QueryBuilder queryBuilder, Class<T> objectClass) {
        SearchRequest request = new SearchRequest(index);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        if (sort != null) {
            Iterator var9 = sort.iterator();

            while (var9.hasNext()) {
                Sort.Order order = (Sort.Order) var9.next();
                sourceBuilder.sort(order.getProperty(), order.getDirection().isDescending() ? SortOrder.DESC : SortOrder.ASC);
            }
        }

        sourceBuilder.fetchSource(fields, null);
        sourceBuilder.query(queryBuilder);
        sourceBuilder.size(scrollSize);
        request.source(sourceBuilder);
        request.scroll(TimeValue.timeValueSeconds(60L));

        SearchResponse searchResponse;
        try {
            searchResponse = this.client.search(request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            throw new RuntimeException("Find scroll document from elasticSearch error.", e);
        }

        String scrollId = searchResponse.getScrollId();
        SearchHits hits = searchResponse.getHits();
        List<T> content = this.parseHits(hits, objectClass);
        return new Scroll(scrollId, new Date((new Date()).getTime() + 60000L), hits.getTotalHits().value, content);
    }

    protected <T> List<T> parseHits(SearchHits hits, Class<T> objectClass) {
        List<T> list = Lists.newArrayList();
        Iterator var4 = hits.iterator();

        while (var4.hasNext()) {
            SearchHit searchHit = (SearchHit) var4.next();

            try {
                T t = objectMapper.readValue(searchHit.getSourceAsString(), objectClass);
                if (t != null) {
                    list.add(t);
                }
            } catch (Exception var7) {
                log.error("Parse searchHit error: " + searchHit.getSourceAsString());
                throw new RuntimeException(var7);
            }
        }

        return list;
    }

    public <T> void fetchAllResults(String index, String sortField, String sortOrder,
                                    Object[] searchAfterValue, Class<T> objectClass,
                                    List<T> results) {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();

        searchSourceBuilder.query(QueryBuilders.matchAllQuery());
        searchSourceBuilder.sort(sortField, SortOrder.fromString(sortOrder));
        searchSourceBuilder.size(10); // 每页返回的文档数

        if (searchAfterValue != null) {
            searchSourceBuilder.searchAfter(searchAfterValue);
        }

        searchRequest.source(searchSourceBuilder);

        SearchResponse searchResponse = null;
        try {
            searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            throw new RuntimeException("Find scroll deep document from elasticSearch error.", e);
        }
        SearchHits hits = searchResponse.getHits();

        // 将当前页的搜索结果转换为目标对象并添加到结果列表中
        List<T> pageResults = parseHits(hits, objectClass);
        results.addAll(pageResults);

        // 获取下一页的 search_after 值
        if (hits.getHits().length > 0) {
            Object[] nextSearchAfterValue = hits.getHits()[hits.getHits().length - 1].getSortValues();
            fetchAllResults(index, sortField, sortOrder, nextSearchAfterValue, objectClass, results);
        }
    }

}
